package com.tsk.base.arthas;

import com.taobao.arthas.agent.attach.ArthasAgent;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.support.DefaultListableBeanFactory;
import org.springframework.context.ApplicationContext;
import org.springframework.jmx.export.annotation.ManagedOperation;
import org.springframework.jmx.export.annotation.ManagedOperationParameter;
import org.springframework.jmx.export.annotation.ManagedResource;
import org.springframework.stereotype.Component;

import java.util.HashMap;
import java.util.Map;

/**
 * @author: 黄嘉振
 * @create: 2023/5/10 16:58
 * @description:    JMX bean
 **/
@Component
@ManagedResource(objectName = "com.ArthasAgentManageMbean:name=ArthasMbean", description = "Arthas远程管理Mbean")
public class ArthasBean {

    @Autowired
    private Map<String, String> arthasConfigMap;

    @Autowired
    private ArthasProperties arthasProperties;

    @Autowired
    private ApplicationContext applicationContext;

    /**
     * 初始化
     *
     * @return
     */
    private ArthasAgent arthasAgentInit() {
        arthasConfigMap = StringUtils.removeDashKey(arthasConfigMap);
        // 给配置全加上前缀
        Map<String, String> mapWithPrefix = new HashMap<String, String>(arthasConfigMap.size());
        for (Map.Entry<String, String> entry : arthasConfigMap.entrySet()) {
            mapWithPrefix.put("arthas." + entry.getKey(), entry.getValue());
        }
        final ArthasAgent arthasAgent = new ArthasAgent(mapWithPrefix, arthasProperties.getHome(),
                arthasProperties.isSlientInit(), null);
        arthasAgent.init();
        return arthasAgent;
    }

    @ManagedOperation(description = "获取配置Arthas Tunnel Server地址")
    public String getArthasTunnelServerUrl() {
        return arthasProperties.getTunnelServer();
    }

    @ManagedOperation(description = "设置Arthas Tunnel Server地址，重新attach后生效")
    @ManagedOperationParameter(name = "tunnelServer", description = "example:ws://127.0.0.1:7777/ws")
    public Boolean setArthasTunnelServerUrl(String tunnelServer) {
        if (tunnelServer == null || tunnelServer.trim().equals("") || tunnelServer.indexOf("ws://") < 0) {
            return false;
        }
        arthasProperties.setTunnelServer(tunnelServer);
        return true;
    }

    @ManagedOperation(description = "获取AgentID")
    public String getAgentId() {
        return arthasProperties.getAgentId();
    }

    @ManagedOperation(description = "获取应用名称")
    public String getAppName() {
        return arthasProperties.getAppName();
    }

    @ManagedOperation(description = "获取ArthasConfigMap")
    public HashMap<String, String> getArthasConfigMap() {
        return (HashMap) arthasConfigMap;
    }

    @ManagedOperation(description = "返回是否已经加载Arthas agent")
    public Boolean isArthasAttched() {
        DefaultListableBeanFactory defaultListableBeanFactory = (DefaultListableBeanFactory) applicationContext.getAutowireCapableBeanFactory();
        String bean = "arthasAgent";
        if (defaultListableBeanFactory.containsBean(bean)) {
            return true;
        }
        return false;
    }

    @ManagedOperation(description = "启动Arthas agent")
    public Boolean startArthasAgent() {
        DefaultListableBeanFactory defaultListableBeanFactory = (DefaultListableBeanFactory) applicationContext.getAutowireCapableBeanFactory();
        String bean = "arthasAgent";
        if (defaultListableBeanFactory.containsBean(bean)) {
            ((ArthasAgent) defaultListableBeanFactory.getBean(bean)).init();
            return true;
        }
        defaultListableBeanFactory.registerSingleton(bean, arthasAgentInit());
        return true;
    }

    @ManagedOperation(description = "关闭Arthas agent，暂未实现")
    public Boolean stopArthasAgent() {
        // TODO 无法获取自定义tmp文件夹加载的classLoader，因此无法获取到com.taobao.arthas.core.server.ArthasBootstrap类并调用destroy方法
        DefaultListableBeanFactory defaultListableBeanFactory = (DefaultListableBeanFactory) applicationContext.getAutowireCapableBeanFactory();
        String bean = "arthasAgent";
        if (defaultListableBeanFactory.containsBean(bean)) {
            defaultListableBeanFactory.destroySingleton(bean);
            return true;
        } else {
            return false;
        }
    }
}
